package com.iteaj.framework.utils;

import com.iteaj.framework.BaseEntity;
import com.iteaj.framework.TreeEntity;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.stream.Collectors;

public class TreeUtils {

    public static  <T extends TreeEntity> Collection<T>
                    toTrees(Collection<T> list, Long parent) {

        if(CollectionUtils.isEmpty(list) || parent == null) return Collections.EMPTY_LIST;

        Map<String, T> menuMap = list.stream().collect(Collectors
                .toMap(item -> item.getId().toString(), item -> item));

        //
        list.forEach(item->{
            T parentItem = menuMap.get(item.getPid().toString());
            if(null != parentItem) parentItem.addChild(item);
        });

        List<T> collect = list.stream().filter(item -> Long
                .compare(item.getPid(), parent) == 0).collect(Collectors.toList());
        return CollectionUtils.isEmpty(collect) ? list : collect;
    }

    public static <T extends TreeEntity> Collection<T> toAdminMenuTrees(List<T> allMenus, List<String> adminMenuIds) {
        Map<String, T> menuMap = allMenus.stream().collect(Collectors
                .toMap(item -> item.getId().toString(), item -> item));

        //
        allMenus.forEach(item->{
            T parentItem = menuMap.get(item.getPid().toString());
            if(null != parentItem) parentItem.addChild(item);
        });

        Map<String, T> adminMenus = new HashMap<>(); // 存放管理员菜单的列表
        adminMenuIds.forEach(item -> {
            T t = menuMap.get(item);
            if(t != null && !adminMenus.containsKey(item)) {
                adminMenus.put(item, t);
            }
            getParentMenu(t.getPid(), menuMap, adminMenus);
        });


        return toTrees(adminMenus.values(), 0l);
    }

    private static <T extends TreeEntity> void getParentMenu(Long pid, Map<String, T> menuMap, Map<String, T> adminMenus) {
        if(pid == null) return;
        String pidStr = String.valueOf(pid);
        T t = menuMap.get(pidStr);
        if(t != null && !adminMenus.containsKey(pidStr)) {
            t.setChildren(null);
            adminMenus.put(pidStr, t);
            getParentMenu(t.getPid(), menuMap, adminMenus);
        }
    }
}
